from win32com.client import Dispatch
import pickle
import argparse

font_size_initial = 8
font_name_initial = "SimHei (TrueType)"

ityp_catRectangle = 1
circle_type = 4
filter_string = "branch_point"

dic_2 = {}


def CreateLine(p1_name, p2_name):
    global dic_point, dic_2, fac2d
    coord1 = dic_point[p1_name]
    coord2 = dic_point[p2_name]
    p1 = dic_2[p1_name]
    p2 = dic_2[p2_name]
    line_ = fac2d.CreateLine(coord1[0], coord1[1], coord2[0], coord2[1])
    line_.StartPoint = p1
    line_.EndPoint = p2
    v_1 = (coord1[0] + coord2[0] ) / 2
    v_2 = (coord1[1] + coord2[1] ) / 2    
    return line_

def CreateLineAssociativeText(line, weight_val, pos_x, pos_y):
    global texts_collection
    text_ = texts_collection.Add(str(weight_val), pos_x, pos_y)
    text_.SetFontSize(0, 0, font_size_initial)
    text_.SetFontName(0,  0, font_name_initial)
    text_.AssociativeElement = line


def layout():
    global dic_point, dic_point_label, dic_line, dic_line_label, dic_2
    for k, v in dic_point.items():
        p = fac2d.CreatePoint(v[0], v[1])
        dic_2[k] = p        
        if dic_point_label.get(k) is not None and dic_point_label[k][0].find(filter_string) < 0:
            point_label_text = dic_point_label[k][0]
            point_label_x = dic_point_label[k][1]
            point_label_y = dic_point_label[k][2]
            text_ = texts_collection.Add(point_label_text, point_label_x, point_label_y)
            text_.SetFontSize(0, 0, font_size_initial)
            text_.SetFontName(0,  0, font_name_initial)
            text_.ActivateFrame(ityp_catRectangle)
            text_.AssociativeElement = p
        selection.Clear()
        selection.Add(p)
        selection.VisProperties.SetRealColor(255,0,0,0)
        selection.VisProperties.SetSymbolType(circle_type)
    for k, v in dic_line.items():
        line_name = k
        point_1_name = v[0]
        point_2_name = v[1]
        weight = v[2]
        line_label_pos_x, line_label_pos_y = dic_line_label[line_name]
        line_ = CreateLine(point_1_name, point_2_name)
        CreateLineAssociativeText(line_, weight, line_label_pos_x, line_label_pos_y)
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--load_pickle", default="0", type=str)
    args = parser.parse_args()
    CATIA = Dispatch("CATIA.Application")
    drawingDoc = CATIA.ActiveDocument
    sheet = drawingDoc.Sheets.ActiveSheet
    View = sheet.Views.ActiveView
    texts_collection = View.Texts
    fac2d = View.Factory2D
    elements = View.GeometricElements
    selection = CATIA.ActiveDocument.Selection
    if args.load_pickle == "0":
        dic_point = {"branch_point": [110, 100], "tx101": [-100, -100], "tx102": [-200, 100], "tx103": [230, 230]}
        dic_edge = [["branch_point", "tx101", 103], ["branch_point", "tx102", 104], ["branch_point", "tx103", 177]]

        dic_point_label = {}
        dic_line_label = {}
        dic_line = {}
        for k, v in dic_point.items():
            dic_point_label[k] = [k, v[0], v[1]]
        for item in dic_edge:
            name_ = item[0] + "_" + item[1]
            dic_line[name_] = item
        for k, v in dic_line.items():
            p1 = dic_point[v[0]]
            p2 = dic_point[v[1]]
            dic_line_label[k] = [(p1[0] + p2[0])/2, (p1[1] + p2[1])/2]
    else:
        with open("graph.pickle", "rb") as f:
            dic = pickle.load(f)
            dic_point = dic["dic_point"]
            dic_point_label = dic["dic_point_label"]
            dic_line = dic["dic_line"]
            dic_line_label = dic["dic_line_label"]
    layout()